package cn.infrabase.crypto.digest;

import cn.infrabase.core.io.FileUtils;
import cn.infrabase.core.io.IoRuntimeException;
import cn.infrabase.core.io.IoUtils;
import cn.infrabase.core.lang.ArrayUtils;
import cn.infrabase.core.util.CharsetUtils;
import cn.infrabase.core.util.HexUtils;
import cn.infrabase.core.util.StringUtils;
import cn.infrabase.crypto.SecurityUtils;

import javax.annotation.Nullable;
import java.io.File;
import java.io.IOException;
import java.io.InputStream;
import java.nio.charset.Charset;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.security.Provider;

/**
 * 摘要器。
 *
 * @author Weimin Gao
 * @since 1.0.0
 **/
public class Digester {

	/**
	 * 消息摘要对象。
	 */
	private MessageDigest digest;

	/**
	 * 盐值。
	 */
	protected byte[] salt;

	/**
	 * 加盐位置，即将盐值放置待散列数据的位置，默认0即放在前面。
	 */
	protected int saltPosition = 0;

	/**
	 * 散列次数，默认1。
	 */
	protected int digestCount = 1;

	/**
	 * 构造函数。
	 *
	 * @param algorithm 摘要算法枚举。
	 * @param provider  摘要算法提供者，如果为null时表示Jdk默认摘要算法。
	 */
	public Digester(DigestAlgorithm algorithm, @Nullable Provider provider) {
		init(algorithm, provider);
	}

	/**
	 * 构造函数。
	 *
	 * @param algorithm 摘要算法，如：MD5。
	 * @param provider  摘要算法提供者，如果为null时表示Jdk默认摘要算法。
	 */
	public Digester(String algorithm, @Nullable Provider provider) {
		this(DigestAlgorithm.codeOf(algorithm), provider);
	}

	/**
	 * 构造函数，使用JDK摘要算法。
	 *
	 * @param algorithm 摘要算法，如：MD5。
	 */
	public Digester(String algorithm) {
		this(algorithm, null);
	}

	/**
	 * 构造函数，使用JDK摘要算法。
	 *
	 * @param algorithm 摘要算法枚举。
	 */
	public Digester(DigestAlgorithm algorithm) {
		this(algorithm, null);
	}

	/**
	 * 初始化摘要器。
	 *
	 * @param algorithm 摘要算法枚举。
	 * @param provider  摘要算法提供者，如果为null时表示使用Jdk默认摘要算法，可以引入其他摘要算法提供者。
	 * @throws DigestAlgorithmNotSupportException 如果摘要算法不支持时。
	 */
	public void init(DigestAlgorithm algorithm, @Nullable Provider provider) throws DigestAlgorithmNotSupportException {
		String code = algorithm.code();
		if (null == provider) {
			this.digest = SecurityUtils.createMessageDigest(algorithm);
		} else {
			try {
				this.digest = MessageDigest.getInstance(code, provider);
			} catch (NoSuchAlgorithmException e) {
				throw new DigestAlgorithmNotSupportException(code);
			}
		}
	}

	/**
	 * 设置加盐内容（默认空）。
	 *
	 * @param salt 盐值。
	 * @return 当前对象。
	 */
	public Digester setSalt(@Nullable byte[] salt) {
		this.salt = salt;
		return this;
	}

	/**
	 * <pre>
	 * 设置加盐的位置，只有盐值存在时有效，加盐的位置指盐位于数据byte数组中的位置。
	 * 例如
	 * 1：当data=0123456时，saltPosition=2时，盐位于data的1和2中间，即：01[salt]23456
	 * </pre>
	 *
	 * @param saltPosition 盐的位置。
	 * @return 当前对象。
	 */
	public Digester setSaltPosition(int saltPosition) {
		this.saltPosition = saltPosition;
		return this;
	}

	/**
	 * 设置摘要计算重复次数。
	 *
	 * @param digestCount 摘要计算重复次数。
	 * @return 当前对象。
	 */
	public Digester setDigestCount(int digestCount) {
		this.digestCount = digestCount;
		return this;
	}

	/**
	 * 重置摘要。
	 *
	 * @return 当前对象。
	 */
	public Digester reset() {
		this.digest.reset();
		return this;
	}

	// -----------------------------------------------------------------------------------------------------------------
	// digest

	/**
	 * 生成摘要。
	 *
	 * @param data 被摘要数据。
	 * @return 摘要数据。
	 */
	public byte[] digest(@Nullable byte[] data) {
		byte[] result;
		if (data == null) {
			result = doDigest((byte[]) null);
		} else {
			if (this.saltPosition <= 0) {
				// 加盐在开头，自动忽略空盐值
				result = doDigest(this.salt, data);
			} else if (this.saltPosition >= data.length) {
				// 加盐在末尾，自动忽略空盐值
				result = doDigest(data, this.salt);
			} else if (ArrayUtils.isNotEmpty(this.salt)) {
				// 加盐在中间
				this.digest.update(data, 0, this.saltPosition);
				this.digest.update(this.salt);
				this.digest.update(data, this.saltPosition, data.length - this.saltPosition);
				result = this.digest.digest();
			} else {
				// 无加盐
				result = doDigest(data);
			}
		}
		return resetAndRepeatDigest(result);
	}

	/**
	 * 生成摘要。
	 *
	 * @param data    被摘要数据。
	 * @param charset 被摘要数据的字符集。
	 * @return 摘要数据。
	 */
	public byte[] digest(String data, @Nullable String charset) {
		return digest(data, CharsetUtils.charset(charset));
	}

	/**
	 * 生成摘要。
	 *
	 * @param data    被摘要数据。
	 * @param charset 被摘要数据的字符集。
	 * @return 摘要数据。
	 */
	public byte[] digest(String data, @Nullable Charset charset) {
		return digest(StringUtils.bytes(data, charset));
	}

	/**
	 * 生成摘要。
	 *
	 * @param data   被摘要数据流。
	 * @param buffer 读取数据流中数据时的缓冲大小。
	 * @return 摘要数据。
	 */
	public byte[] digest(InputStream data, int buffer) {
		if (buffer < 1) {
			buffer = IoUtils.DEFAULT_BUFFER_SIZE;
		}
		byte[] result;
		try {
			if (ArrayUtils.isEmpty(this.salt)) {
				result = digestWithoutSalt(data, buffer);
			} else {
				result = digestWithSalt(data, buffer);
			}
		} catch (IOException e) {
			throw new IoRuntimeException(e, String.format("生成摘要发生错误【%s】", e.getMessage()));
		}
		return resetAndRepeatDigest(result);
	}

	/**
	 * 生成摘要。
	 *
	 * @param data 被摘要数据流。
	 * @return 摘要数据。
	 */
	public byte[] digest(InputStream data) {
		return digest(data, IoUtils.DEFAULT_BUFFER_SIZE);
	}

	/**
	 * 生成文件摘要。
	 *
	 * @param file 被摘要的文件。
	 * @return 摘要数据。
	 */
	public byte[] digest(File file) {
		InputStream is = null;
		try {
			is = FileUtils.getInputStream(file);
			return digest(is);
		} finally {
			IoUtils.close(is);
		}
	}

	// -----------------------------------------------------------------------------------------------------------------
	// digestHex

	/**
	 * 生成摘要，并转为16进制字符串。
	 *
	 * @param data 被摘要数据。
	 * @return 摘要数据。
	 */
	public String digestHex(byte[] data) {
		return HexUtils.encodeHexStr(digest(data));
	}

	/**
	 * 生成摘要，并转为16进制字符串。
	 *
	 * @param data    被摘要数据。
	 * @param charset 被摘要数据字符集。
	 * @return 摘要数据。
	 */
	public String digestHex(String data, @Nullable Charset charset) {
		return HexUtils.encodeHexStr(digest(data, charset));
	}

	/**
	 * 生成摘要，并转为16进制字符串。
	 *
	 * @param data    被摘要数据。
	 * @param charset 被摘要数据字符集，如：UTF-8。
	 * @return 摘要数据。
	 */
	public String digestHex(String data, @Nullable String charset) {
		return digestHex(data, CharsetUtils.charset(charset));
	}

	/**
	 * 生成摘要，并转为16进制字符串。
	 *
	 * @param data 被摘要数据，字符集为UTF-8。
	 * @return 摘要。
	 */
	public String digestHex(String data) {
		return digestHex(data, CharsetUtils.CHARSET_UTF_8);
	}

	/**
	 * 生成摘要，并转为16进制字符串。
	 *
	 * @param data 被摘要数据。
	 * @return 摘要数据。
	 */
	public String digestHex(InputStream data) {
		return HexUtils.encodeHexStr(digest(data));
	}

	/**
	 * 生成摘要，并转为16进制字符串。
	 *
	 * @param data   被摘要数据。
	 * @param buffer 缓存大小，小于1时使用默认缓存大小（{@link IoUtils#DEFAULT_BUFFER_SIZE}）。
	 * @return 摘要数据。
	 */
	public String digestHex(InputStream data, int buffer) {
		return HexUtils.encodeHexStr(digest(data, buffer));
	}

	/**
	 * 生成文件摘要，并转为16进制字符串。
	 *
	 * @param file 被摘要的文件。
	 * @return 摘要数据。
	 */
	public String digestHex(File file) {
		return HexUtils.encodeHexStr(digest(file));
	}

	// -----------------------------------------------------------------------------------------------------------------

	/**
	 * 生成摘要。
	 *
	 * @param data 待摘要的数据。
	 * @return 摘要结果。
	 */
	private byte[] doDigest(byte[]... data) {
		for (byte[] value : data) {
			if (null != value) {
				this.digest.update(value);
			}
		}
		return this.digest.digest();
	}

	/**
	 * <pre>
	 * 重复计算摘要，取决于{@link #digestCount}值。
	 * 每次计算摘要前都会重置{@link #digest}。
	 * </pre>
	 *
	 * @param digestData 第一次摘要过的数据。
	 * @return 摘要结果。
	 */
	private byte[] resetAndRepeatDigest(byte[] digestData) {
		final int digestCount = Math.max(1, this.digestCount);
		reset();
		for (int i = 0; i < digestCount - 1; i++) {
			digestData = doDigest(digestData);
			reset();
		}
		return digestData;
	}

	/**
	 * 生成摘要。
	 *
	 * @param data   被摘要的数据流。
	 * @param buffer 缓存大小，小于1时使用默认缓存大小（{@link IoUtils#DEFAULT_BUFFER_SIZE}）。
	 * @return 摘要结果。
	 */
	private byte[] digestWithoutSalt(InputStream data, int buffer) throws IOException {
		final byte[] bytes = new byte[buffer];
		int read;
		while ((read = data.read(bytes, 0, buffer)) > -1) {
			this.digest.update(bytes, 0, read);
		}
		return this.digest.digest();
	}

	/**
	 * 生成摘要。
	 *
	 * @param data   被摘要的数据流。
	 * @param buffer 缓存大小，小于1时使用默认缓存大小（{@link IoUtils#DEFAULT_BUFFER_SIZE}）。
	 * @return 摘要结果。
	 */
	private byte[] digestWithSalt(InputStream data, int buffer) throws IOException {
		if (this.saltPosition <= 0) {
			// 加盐在开头
			this.digest.update(this.salt);
		}
		final byte[] bytes = new byte[buffer];
		int total = 0;
		int read;
		while ((read = data.read(bytes, 0, buffer)) > -1) {
			total += read;
			if (this.saltPosition > 0 && total >= this.saltPosition) {
				if (total != this.saltPosition) {
					digest.update(bytes, 0, total - this.saltPosition);
				}
				// 加盐在中间
				this.digest.update(this.salt);
				this.digest.update(bytes, total - this.saltPosition, read);
			} else {
				this.digest.update(bytes, 0, read);
			}
		}
		if (total < this.saltPosition) {
			// 加盐在末尾
			this.digest.update(this.salt);
		}
		return this.digest.digest();
	}
}
